#include "ir_constOptimizer.h"
#include <unordered_map>
#include <vector>
#include <variant>

namespace ir{
    bool constOptimizer::PropagateAndFoldConstants(Function& func){
        std::unordered_map<std::string, int> constMap;
        std::vector<std::pair<Block*, std::vector<Instruction>::iterator>> toRemove;//待删除指令的信息（块、指令位置）
        bool changed = false;
        for(auto& block : func.blocks){
            std::unordered_map<std::string, int> stackMap; // 当前基本块的栈变量映射
            for(auto it = block.instrs.begin();it != block.instrs.end();++it){
                //常量传播
                auto& instr = *it;
            
                
                //STORE指令  store value, @mem
                if(instr.op == OP_STORE){
                    if (instr.srcs.size() == 2) {
                        if (auto* mem_var = std::get_if<std::string>(&instr.srcs[1])) {
                            if (std::holds_alternative<int>(instr.srcs[0])) {
                                // 存储常量值
                                stackMap[*mem_var] = std::get<int>(instr.srcs[0]);
                                //std::cout<<* mem_var<<" "<<std::get<int>(instr.srcs[0])<<"\n";
                            } else {
                                // 存储非常量，移除已知值
                                if(stackMap.find(*mem_var) != stackMap.end()){
                                    stackMap.erase(*mem_var);
                                    //std::cout<<"FOUND:"<<*mem_var;
                                }
                                else{
                                    //std::cout<<"NOT FOUND\n";
                                }
                            }
                        }
                    }
                }
                
                //LOAD指令 %dest = load @mem
                if(instr.op == OP_LOAD){
                    if (auto* mem_var = std::get_if<std::string>(&instr.srcs[0])) {
                        // 如果栈变量有已知常量值
                        if (stackMap.find(*mem_var) != stackMap.end()) {
                            constMap[instr.dest] = stackMap[*mem_var];
                            toRemove.push_back({&block, it});
                            changed = true;
                        }
                    }
                }
                
                //临时常量
                for(auto& src : instr.srcs){
                    if(auto* var = std::get_if<std::string>(&src)){
                        if(constMap.find(*var) != constMap.end()){
                            src = constMap[*var];
                            changed = true;
                        }
                    }
                }
                

                //常量折叠
                if(instr.op >= OP_ADD && instr.op <=OP_SAR){//所有的二元运算
                    if(instr.srcs.size() == 2 && std::holds_alternative<int>(instr.srcs[0]) && std::holds_alternative<int>(instr.srcs[1])){
                            int left = std::get<int>(instr.srcs[0]);
                            int right = std::get<int>(instr.srcs[1]);
                            int res = ComputeBinaryOp(instr.op, left, right);
                            constMap[instr.dest] = res;
                            toRemove.push_back({&block, it});
                            changed = true;
                        }
                }
                /*
                if(instr.op == OP_BR){
                    if (instr.srcs.size() == 3 && std::holds_alternative<int>(instr.srcs[0])) 
                    {
                        int cond = std::get<int>(instr.srcs[0]);
                        // 转换为无条件跳转
                        instr.op = OP_JMP;
                        instr.srcs = {cond ? instr.srcs[1] : instr.srcs[2]};
                        changed = true;
                    }
                }
                */
            }
        }
        //删除被折叠指令
        for(auto it = toRemove.rbegin(); it != toRemove.rend(); ++it){
            auto&[block, instrIt] = *it;
            block -> instrs.erase(instrIt);
        }
        return changed;
    }
    int constOptimizer::ComputeBinaryOp(Opcode op, int left, int right) {
        switch (op) {
            case OP_ADD: return left + right;
            case OP_SUB: return left - right;
            case OP_MUL: return left * right;
            case OP_DIV: return right != 0 ? left / right : 0;
            case OP_MOD: return right != 0 ? left % right : 0;
            case OP_EQ:  return left == right;
            case OP_NE:  return left != right;
            case OP_GT:  return left > right;
            case OP_LT:  return left < right;
            case OP_GE:  return left >= right;
            case OP_LE:  return left <= right;
            case OP_AND: return left && right;
            case OP_OR:  return left || right;
            case OP_XOR: return left ^ right;
            case OP_SHL: return left << right;
            case OP_SHR: return left >> right;
            case OP_SAR: return left >> right;
            default: return 0;
        }
    }
}